package org.simplesocket.simplenettycore.handlers;

import cn.hutool.core.lang.UUID;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelDuplexHandler;
import io.netty.channel.ChannelHandlerContext;
import lombok.extern.slf4j.Slf4j;
import org.simplesocket.simplenettycore.context.SimpleNettyContext;
import org.simplesocket.simplenettycore.protocol.SimpleProtocol;
import org.simplesocket.simplenettycore.session.ChannelAttr;
import org.simplesocket.simplenettycore.session.ClientSessionContext;
import org.simplesocket.simplenettycore.utils.CommonUtils;
import org.simplesocket.simplenettycore.utils.SimpleNettyUtils;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

@Slf4j
public class ConnectLifecycleHandler extends ChannelDuplexHandler {
    
    /**
     * 加载一些回调函数，在本库中，有很多的pipeline，pipeline中定义了许多的
     * 出站入站处理器，有些时候库使用者需要在某些出站入站处理器调用时调用一些
     * 库使用者自定义的回调函数，例如ConnectLifecycleHandler，管理连接生
     * 命周期的处理器，它有 开始连接、连接建立、连接断开、断线重连 等事件，
     * 在这些事件发生的时候，库使用者或许会希望做一些操作，于是在ConnectLifecycleHandler
     * 处理器中我们定义了一个成员变量（接口），它有一些方法，用户只要实现了这个
     * 接口，并且在接口上添加@Callback注解，那么程序启动时就会去找接口实现类，
     * 实例化它并将它的实例化对象注入到这个成员变量的值中，那么在ConnectLifecycleHandler
     * 中对成员变量（接口）的方法调用就会调用到库使用者自定义的方法上，从而完成回调
     * 功能。这里的原理是：
     * <p>
     * pipeline.addLast("",ChannelPipelineInitializer.connectLifecycleHandler)
     * 但是ChannelPipelineInitializer.connectLifecycleHandler的值还未初始化，
     * 需要在程序启动时初始化它，在ContextHelper.loadCallbacks()方法中我们做了初始化
     */
    private final ClientEventCallback clientEventCallback;
    
    private final SimpleNettyContext context;

    public ConnectLifecycleHandler(SimpleNettyContext context,ClientEventCallback clientEventCallback) {
        this.clientEventCallback = clientEventCallback;
        this.context = context;
    }
    
    @Override
    public void channelActive(ChannelHandlerContext ctx) {
        Channel channel = ctx.channel();
        try {
            // 客户端第一次连接时触发该方法，我们在这里处理初次连接
            // 生成一个客户端唯一值，并将值放入客户端channel中
            String sessionID = UUID.randomUUID().toString(true);
            ClientSessionContext csc = new ClientSessionContext();
            csc.setSessionID(sessionID);
            channel.attr(ChannelAttr.CLIENT_SESSION_CONTEXT).set(csc);
            
            // 将生成的客户端唯一值放入一个HashSet中，方便以后判断新连接的请求是否合法
            context.getRegisteredClient().add(sessionID);
            
            // 将数据写回客户端
            ProtocolDataWriter.writeDataBack(ctx, SimpleProtocol.builder()
                    .type(SimpleProtocol.Type.CONNECTED)
                    .status(SimpleProtocol.Status.SUCCESS)
                    .sessionID(sessionID)
                    .data(JSON.toJSONString(Map.of("sessionID",sessionID)))
                    .build());
            
            log.info("A client try to connect server,remote address is '{}', session id is '{}'",
                    channel.remoteAddress(), sessionID);
            
            // 客户端第一次连接的回调，这里回调的是库使用者实现的ClientEventHandler，
            // 库使用者在实现ClientEventHandler时还需要在类上加注解@ClientEventListener
            if (clientEventCallback != null) {
                context.getSimpleNettyConfig().getWorker()
                        .submit(() -> clientEventCallback.onFirstConnect(sessionID));
            }
        }catch (Exception e){
            log.error("A client '{}' connect server error:\n",channel.remoteAddress());
            CommonUtils.printException(e);
            
            ProtocolDataWriter.writeDataBack(ctx,SimpleProtocol.builder()
                            .type(SimpleProtocol.Type.CONNECTED)
                            .status(SimpleProtocol.Status.ERROR)
                            .data(JSON.toJSONString(Map.of("message","Connect server error")))
                    .build());
        }
    }
    
    @Override
    public void channelInactive(ChannelHandlerContext ctx) {
        // 客户端断开连接时触发这个方法，在这里我们统一做处理
        // 断开连接时清除存储的客户端唯一值
        Channel channel = ctx.channel();
        try {
            ClientSessionContext csc = channel.attr(ChannelAttr.CLIENT_SESSION_CONTEXT).get();
            if (csc != null){
                String sessionID = csc.getSessionID();
                context.getRegisteredClient().remove(sessionID);
                context.getConnectedClient().remove(sessionID);

                log.info("A client with address '{}' disconnected ...", channel.remoteAddress());

                // 客户端断开连接的回调，传递一个ClientSessionContext对象
                if (clientEventCallback != null) {
                    context.getSimpleNettyConfig().getWorker().submit(() -> clientEventCallback.onDisconnect(csc));
                }
            }
        }catch (Exception e){
            log.error("A client '{}' disconnect error:\n",channel.remoteAddress());
            CommonUtils.printException(e);
        }finally {
            ctx.close();
        }
    }
    
    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) {
        ByteBuf buf = (ByteBuf) msg;

        buf.readInt();// 去除长度的4字节
        int magicNumber = buf.readInt(); // 读取开头魔数占的4字节
        if (magicNumber != SimpleProtocol.MAGIC_NUMBER) {
            // 魔数不正确，丢弃数据关闭连接
            ctx.close();
            return;
        }
        
        // 读取请求类型
        byte type = buf.readByte();
        // 读取状态
        buf.readByte();
        
        // 判断sessionID，没有该sessionID的话（除了重连类型）是非法连接，直接关闭
        String sessionID = buf.readBytes(32).toString(StandardCharsets.UTF_8);
        if (StrUtil.isBlank(sessionID)
                || !(context.getRegisteredClient().contains(sessionID)
                || type == SimpleProtocol.Type.RECONNECT.getValue())
        ) {
            ctx.close();
            return;
        }
        
        if (type == SimpleProtocol.Type.NORMAL.getValue()) {
            // 普通类型，直接放行
            ctx.fireChannelRead(buf.resetReaderIndex());
        } else {
            // 特殊类型，特殊处理
            handleSpecialType(ctx, buf.resetReaderIndex());
        }
    }
    
    private void handleSpecialType(ChannelHandlerContext ctx, ByteBuf buf) {
        buf.readInt(); // 去除开头长度占的4字节
        buf.readInt(); // 去除开头魔数占的4字节
        byte type = buf.readByte();// 获取类型1字节
        buf.readByte();// 去除状态
        String sessionID = buf.readBytes(32).toString(StandardCharsets.UTF_8);// 读取sessionID占的32字节
        int uriLength = buf.readInt();// 读取uri长度
        buf.readBytes(uriLength);// 读取uri
        String data = buf.readBytes(buf.readableBytes()).toString(StandardCharsets.UTF_8);// 读取数据剩余所有字节
        ClientSessionContext csc = ctx.channel().attr(ChannelAttr.CLIENT_SESSION_CONTEXT).get();
        
        switch (type) {
            case 1 -> ctx.close();// 客户端不会发这个类型的连接，直接关闭连接
            case 2 -> handleGetKeys(ctx, data, csc);
            case 3 -> handleReconnect(ctx, data, csc, sessionID);
            default -> ctx.close();// 没有这种类型的连接，直接关闭连接丢弃数据
        }
    }
    
    private void handleGetKeys(ChannelHandlerContext ctx, String clientPubRsa, ClientSessionContext csc) {
        // 客户端想要获取加密key，但是自身得携带一个公钥key
        try {
            String coarseAES = SimpleNettyUtils.AESUtil.generateRandomKeyString();// 生成aes
            List<String> keyStrings = SimpleNettyUtils.RSAUtil.generateKeyStrings();// 生成RSA公私钥
            String serverPubRSA = keyStrings.getFirst();// 服务器公钥
            String serverPriRSA = keyStrings.getLast();// 服务器私钥
            
            String encryptedAES = SimpleNettyUtils.RSAUtil.encrypt(clientPubRsa, coarseAES);// 加密AES密钥
            
            // 从channel上下文中获取数据存储对象
            csc.setBothAes(coarseAES);
            csc.setClientPubRsa(clientPubRsa);
            csc.setServerPubRsa(serverPubRSA);
            csc.setServerPriRsa(serverPriRSA);
            ctx.channel().attr(ChannelAttr.CLIENT_SESSION_CONTEXT).set(csc);// 将数据写回channel上下文
            
            record Keys(String serverPubRSA, String aes) {
            }
            Keys keys = new Keys(serverPubRSA, encryptedAES);// 包装keys
            
            // 数据写回客户端
            ProtocolDataWriter.writeDataBack(ctx, SimpleProtocol.builder()
                            .type(SimpleProtocol.Type.DISCUSS_KEYS)
                            .status(SimpleProtocol.Status.SUCCESS)
                            .sessionID(csc.getSessionID())
                            .data(JSON.toJSONString(keys))
                    .build());
            
            // 回调库使用者的方法
            if (clientEventCallback != null) {
                context.getSimpleNettyConfig().getWorker().submit(() -> clientEventCallback.onConnected(csc));
            }
            
        } catch (Exception e) {
            log.error("JSON serialization failed at 'ConnectLifecycleHandler.handleSpecialType()'\n");
            CommonUtils.printException(e);
            
            ProtocolDataWriter.writeDataBack(ctx, SimpleProtocol.builder()
                            .type(SimpleProtocol.Type.DISCUSS_KEYS)
                            .status(SimpleProtocol.Status.ERROR)
                            .sessionID(csc.getSessionID())
                            .data(JSON.toJSONString(Map.of("message","Get keys error")))
                    .build());
        }
    }
    
    private void handleReconnect(ChannelHandlerContext ctx, String previousSessionID,
                                 ClientSessionContext csc, String currentSessionID) {
        try {
            // 重连类型，从data中获取该连接之前的sessionID
            // 获取当前的sessionID并在已注册sessionID表中移除
            context.getRegisteredClient().remove(currentSessionID);// 移除当前的sessionID
            context.getRegisteredClient().add(previousSessionID);// 将之前的sessionID添加
            context.getConnectedClient().remove(currentSessionID);// 移除当前的sessionID与当前channel映射
            context.getConnectedClient().put(previousSessionID, ctx.channel());// 添加之前的sessionID与当前channel映射
            
            // 将ctx中的sessionID替换更新
            csc.setSessionID(previousSessionID);
            ctx.channel().attr(ChannelAttr.CLIENT_SESSION_CONTEXT).set(csc);
            
            ProtocolDataWriter.writeDataBack(ctx, SimpleProtocol.builder()
                    .type(SimpleProtocol.Type.RECONNECT)
                    .status(SimpleProtocol.Status.SUCCESS)
                    .sessionID(JSON.toJSONString(Map.of("previousSessionID",previousSessionID)))
                    .build());
            
            // 重连回调库使用者定义的逻辑
            if (clientEventCallback != null) {
                context.getSimpleNettyConfig().getWorker()
                        .submit(() -> clientEventCallback.onReconnected(previousSessionID));
            }
        }catch (Exception e){
            log.error("A client '{}' reconnect server error:\n",ctx.channel().remoteAddress());
            CommonUtils.printException(e);
            
            ProtocolDataWriter.writeDataBack(ctx,SimpleProtocol.builder()
                            .type(SimpleProtocol.Type.RECONNECT)
                            .status(SimpleProtocol.Status.ERROR)
                            .data(JSON.toJSONString(Map.of("message","Reconnect server error")))
                    .build());
        }
    }
}
